import os

from pymilvus import MilvusClient
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.core import StorageContext
from pymilvus import model

import json

import ollama

MODEL = 'ishumilin/deepseek-r1-coder-tools:1.5b'
MODEL = 'ishumilin/deepseek-r1-coder-tools:14b'
MODEL='qwen2.5-coder:7b'
# MODEL='deepseek-r1:7b'
# 创建一个集合并插入数据

milvus_client = None


def initMilvusClient():
    global milvus_client
    if milvus_client is not None:
        return
    # milvus_client = MilvusClient('milvus_local.db') #理论上相当于运行 milvus-server --data milvus_local.db
    milvus_client = MilvusClient(
        uri="http://localhost:19530",
        token="root:Milvus",

    )
    print('initMilvusClient ')


embedding_fn = model.DefaultEmbeddingFunction()


# 模拟获取航班时间的API调用
# 在实际应用中，这会从实时数据库或API获取数据

def get_flight_times(departure: str, arrival: str) -> str:
    flights = {

        'NYC-LAX': {'departure': '08:00 AM', 'arrival': '11:30 AM', 'duration': '5h 30m'},

        'LAX-NYC': {'departure': '02:00 PM', 'arrival': '10:30 PM', 'duration': '5h 30m'},

        'LHR-JFK': {'departure': '10:00 AM', 'arrival': '01:00 PM', 'duration': '8h 00m'},

        'JFK-LHR': {'departure': '09:00 PM', 'arrival': '09:00 AM', 'duration': '7h 00m'},

        'CDG-DXB': {'departure': '11:00 AM', 'arrival': '08:00 PM', 'duration': '6h 00m'},

        'DXB-CDG': {'departure': '03:00 AM', 'arrival': '07:30 AM', 'duration': '7h 30m'},

    }

    key = f'{departure}-{arrival}'.upper()

    return json.dumps(flights.get(key, {'error': '航班未找到'}))


# 在向量数据库中搜索关于人工智能的信息

def search_data_in_vector_db(query: str) -> str:
    query_vectors = embedding_fn.encode_queries([query])
    initMilvusClient()
    res = milvus_client.search(
        collection_name="demo_collection",
        data=query_vectors,
        limit=2,
        output_fields=["text", "subject"],  # 指定返回哪些字段

    )

    print(res)
    return json.dumps(res)


def test1(question: str):
    initMilvusClient()
    client = ollama.Client(host='http://localhost:11434')

    # 初始化对话，包含用户提问

    messages = [{'role': 'user', 'content': question}]

    # 第一次 API 调用：发送查询和函数描述给模型

    print(f'MODEL={MODEL}')
    response = client.chat(

        model=MODEL,

        messages=messages,

        tools=[

            {

                'type': 'function',

                'function': {

                    'name': 'get_flight_times',

                    'description': '获取两个城市之间的航班时间',

                    'parameters': {

                        'type': 'object',

                        'properties': {

                            'departure': {

                                'type': 'string',

                                'description': '出发城市（机场代码）',

                            },

                            'arrival': {

                                'type': 'string',

                                'description': '目的城市（机场代码）',

                            },

                        },

                        'required': ['departure', 'arrival'],

                    },

                },

            },

            {

                'type': 'function',

                'function': {

                    'name': 'search_data_in_vector_db',

                    'description': '在向量数据库中搜索关于人工智能的数据',

                    'parameters': {

                        'type': 'object',

                        'properties': {

                            'query': {

                                'type': 'string',

                                'description': '搜索查询',

                            },

                        },

                        'required': ['query'],

                    },

                },

            },

        ],

    )

    # 将模型的响应添加到对话历史中

    messages.append(response['message'])

    # 检查模型是否使用了提供的函数

    if not response['message'].get('tool_calls'):
        print("模型没有使用函数。它的响应是：")

        print(response['message']['content'])

        return

        # 处理模型发出的函数调用

    if response['message'].get('tool_calls'):

        available_functions = {

            'get_flight_times': get_flight_times,

            'search_data_in_vector_db': search_data_in_vector_db,

        }

        for tool in response['message']['tool_calls']:
            function_to_call = available_functions[tool['function']['name']]

            function_args = tool['function']['arguments']

            function_response = function_to_call(**function_args)

            # 将函数响应添加到对话中

            messages.append(

                {

                    'role': 'tool',

                    'content': function_response,

                }

            )

            # 第二次 API 调用：获取模型的最终响应

    final_response = client.chat(model=model, messages=messages)

    print(final_response['message']['content'])


def test3():
    # pip install -U huggingface_hub
    # set "HF_ENDPOINT=https://hf-mirror.com"
    # huggingface-cli download --resume-download GPTCache/paraphrase-albert-onnx --local-dir F:\HuggingFaceModels\GPTCache\paraphrase-albert-onnx
    # 运行python 时候增加变量 HF_ENDPOINT=https://hf-mirror.com
    # 下载的模型位于 C:\Users\Administrator\.cache\huggingface\hub
    # 也可以设置默认模型存放目录,增加环境变量 HF_HOME=F:\HuggingFaceModels
    embedding_fn = model.DefaultEmbeddingFunction()

    docs = [

        "人工智能作为一门学科在1956年成立。",

        "艾伦·图灵是最早进行实质性人工智能研究的人之一。",

        "图灵出生在伦敦的梅达维尔，并在英格兰南部长大。",

    ]

    vectors = embedding_fn.encode_documents(docs)

    # 输出的向量有768个维度，正好与我们刚刚创建的集合相匹配。

    print("维度:", embedding_fn.dim, vectors[0].shape)  # 维度: 768 (768,)

    # 每个实体都有id，向量表示，原始文本和主题标签。

    data = [

        {"id": i, "vector": vectors[i], "text": docs[i], "subject": "历史"}

        for i in range(len(vectors))

    ]

    print("数据包含", len(data), "个实体，每个实体包含的字段为：", data[0].keys())

    print("向量维度：", len(data[0]["vector"]))

    milvus_client.create_collection(

        collection_name="demo_collection",

        dimension=768,  # 用于此次演示的向量具有768个维度

    )

    milvus_client.insert(collection_name="demo_collection", data=data)


if __name__ == '__main__':
    question = "从纽约（NYC）到洛杉矶（LAX）的飞行时长是多少？"
    error = True
    i = 2
    while error and i > 0:
        i -= 1
        try:
            test1(question)
            error = False
        except Exception as e:
            print(e)
            continue

    # test2()
    # test3()
